import requests
import json


def access_token():
    api_key = "NAjHiybRP3UXzENY0E4XDsfo"
    secret_key = "W23YOtWRGLHqSSxEsZ1e6Wvwxin1Fy9S"

    url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={api_key}&client_secret={secret_key}"

    payload = ""
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload).json()

    return response["access_token"]


def ernie_three(data):
    """
        ernir_three_8k 接口，接收一个data，是一个dict对象，详细如下：
        data: {
            "message": [] # 列表，
            "config": {
                "temperature": 0-2, default 1，
                "max_token": 0-1024, default 512
            }
        }
    """
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + access_token()

    payload = json.dumps({
        "messages": data["message"][1:],
        "system": data["message"][0]["content"],
        "temperature": data["config"]["temperature"] if (data["config"]["temperature"] >= 0) and (
                data["config"]["temperature"] <= 1) else 0.5,  # 范围是0-1
        "max_output_tokens": data["config"]["max_token"] if (data["config"]["max_token"] >= 2) and (
                data["config"]["max_token"] <= 1024) else 512,  # 范围2-2048
    })

    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload).json()

    print(response)
    print(response["result"])


def ernie_three_stream(data):
    """
        ernir_three_8k 接口，接收一个data，是一个dict对象，详细如下：
        data: {
            "message": [] # 列表，
            "config": {
                "temperature": 0-2, default 1，
                "max_token": 0-1024, default 512
            }
        }
    """
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + access_token()

    payload = json.dumps({
        "messages": data["message"][1:],
        "system": data["message"][0]["content"],
        "temperature": data["config"]["temperature"] if (data["config"]["temperature"] >= 0) and (
                data["config"]["temperature"] <= 1) else 0.5,  # 范围是0-1
        "max_output_tokens": data["config"]["max_token"] if (data["config"]["max_token"] >= 2) and (
                data["config"]["max_token"] <= 1024) else 512,  # 范围2-2048
        "stream": True
    })

    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload, stream=True)

    for line in response.iter_lines():
        if line:
            line = line.decode("UTF-8")
            colon_index = line.find(':')
            json_str = line[colon_index + 1:].strip()
            data = json.loads(json_str)["result"]
            yield data


if __name__ == '__main__':
    data1 = {
        "message": [
            {"role": "system", "content": "你是一个很有帮助的小男孩"},
            {"role": "user", "content": "你好啊"}
        ],
        "config": {
            "temperature": 1,
            "max_token": 512
        }
    }

    ernie_three_stream(data1)
